﻿module Environment

open TinyJava.Ast
open ExceptionMonad
open System.Collections.Generic

type Globals = Dictionary<string,ClassTable>

and GenericTable =
| ClassTable of ClassTable
| BlockTable of BlockTable
| MethodInfo of MethodInfo

and ClassTable =
  {
    Name : Id
    InheritedClass : Option<Id>
    Fields : Dictionary<string,InstanceVarSymbol>
    Methods : Dictionary<string,MethodInfo>
  }

and InstanceVarSymbol =
  {
    Name : string
    Type : SynType
  }

and ConstructorInfo =
  {
    Name : string
    Params : Formals list
    Body : Ref<Option<BlockTable>>
  }

and MethodInfo =
  {
    Name : string
    RetType : MethodRetType
    Params  : Formals list
    Body    : Ref<Option<BlockTable>>
    OwnerClass : Ref<ClassTable>
  }

and BlockTable =
  {
    mutable Vars : (string * SynType) list
    mutable Parent  : Option<Ref<BlockTable>>
    mutable Owner   : Option<Ref<MethodInfo>>
    mutable Child   : Option<Ref<BlockTable>>
  }
  member this.addVar (v : (string * SynType)) =
    this.Vars <- this.Vars @ [v]
  member this.isRoot =
    match this.Parent with
    | Some parent -> false
    | None -> true

and Environment =
  {
    SymbolTable   : Globals            //symbol table
    BlockTable    : Ref<BlockTable>    //current block table
    CurrentClass  : Id            //current class id
    CheckReturn   : bool          //true if we must check return statements
    ReturnFound : bool    //true if a return has been found in a control structure body
    LoopCounter : int   //used to check if we are inside the scope of a loop. 0 means we are not in a loop.
  }

let is_library_function (f_name : Id) : bool =
  match f_name with
  | "println"
  | "print"
  | "printi"
  | "printb"
  | "readi"
  | "readln"
  | "eof"
  | "stoi"
  | "itos"
  | "stoa"
  | "aots"
  | "random"
  | "time"
  | "exit" -> true
  | _ -> false

let try_add (table : Dictionary<'k,'v>) (entry : 'k * 'v) (error : string) =
  let k,v = entry
  if table.ContainsKey(k) then
    failwith error
  else
    table.Add(k,v)

let try_get (table : Dictionary<'k,'v>) (key : 'k) =
  if table.ContainsKey(key) then
    Some (table.[key])
  else
    None
    
let rec check_inherited_field (parent : Id) (table : Globals) (field : Id) (error : string) =
  let class_table = table.[parent]
  if class_table.Fields.ContainsKey(field) then
    failwith error

  match class_table.InheritedClass with
  | Some c -> check_inherited_field c table field error
  | None -> ()

//fill the symbol table with class definitions and methods
let fill_table (ast : Program) : Globals =
  let (table : Globals) = Dictionary<string,ClassTable>()

  //error strings for the dictionary
  let field_error =  "Class field already defined"
  let method_error =  "Method already defined"
  let class_error = "Class already defined"
  let main_not_found_error = "Main method not found"
  let multiple_main_error = "Only one main method can be defined"
  let main_return_error = "Main method does not return a value"
  let main_arg_error = "Main must take a SINGLE argument of type string[] or no arguments"
  let inherited_error = "Inherited class definition missing"
  let inherited_field_error = "Class field already exists in a parent class"

  //pick each class in AST and build an entry for the symbol table
  for class_name,inherited_class,members in ast do
    let field_table = Dictionary<string,InstanceVarSymbol>()
    let method_table = Dictionary<string,MethodInfo>()

    let class_table =
      {
        Name = class_name
        InheritedClass = inherited_class
        Fields = field_table
        Methods = method_table
      }

    match inherited_class with
    | Some c ->
        //inherited class must be defined
        if table.ContainsKey(c) |> not then
          failwith inherited_error

    | None -> ()

    //pick each member of a class and build an entry for the symbol table
    for fm in members do
      match fm with
      | FieldDecl(t,vars) ->
          for v in vars do
            let var_symbol = {Name = v; Type = t}
            //check for inherited fields with the same name
            match inherited_class with
            | Some c -> check_inherited_field c table v inherited_field_error
            | None -> ()
            
            try_add field_table (v,var_symbol) field_error
      | MethodDecl(ret_type,name,formals,_) ->
        match ret_type with
        | RetType(t) ->
            let method_info = 
              {
                RetType = RetType(t)
                Params = formals 
                Body = ref None
                Name = name
                OwnerClass = ref class_table
              }
            try_add method_table (name,method_info) method_error
        | Void ->
            let method_info = 
              {
                RetType = Void
                Params = formals
                Body = ref None
                Name = name
                OwnerClass = ref class_table
              }
            try_add method_table (name,method_info) method_error
    try_add table (class_name,class_table) class_error

  //we checke that main exists and its return type is VOID
  let method_list =
    [for c in table do
      yield! c.Value.Methods]

  //check if main is defined
  let main_existing = method_list |> Seq.tryFind(fun m -> m.Key = "main")
  match main_existing with
  | None -> failwith main_not_found_error
  | Some _ -> ()

  //check if main is unique
  let main_list =
    [for m in method_list do
      if m.Key = "main" then
        yield m]
  if main_list |> List.length > 1 then
    failwith multiple_main_error

  //check if main has a correct return type
  let main_correct_return = method_list |> Seq.tryFind(fun m -> m.Key = "main" && m.Value.RetType = Void)
  match main_correct_return with
  | None -> failwith main_return_error
  | Some _ -> ()

  //check if main has a correct argument
  let main_correct_arg = method_list |> Seq.tryFind(fun m ->  m.Key = "main" && 
                                                              m.Value.RetType = Void && 
                                                              ((m.Value.Params |> List.length = 1 &&
                                                              m.Value.Params |> List.exists(fun (pt,_) -> pt = Array(String)) || 
                                                              m.Value.Params |> List.length = 0)))
  match main_correct_arg with
  | None -> failwith main_arg_error
  | Some _ -> ()
  
  //main is correct and we return the symbol table
  table

let rec try_find_id (table : GenericTable) (id : Id) (check_block_only : bool) (env : Environment) : Option<SynType> =
  match table with
  | BlockTable(table) ->
    let method_var_opt = table.Vars |> List.tryFind(fun (n,t) -> n = id)
    match method_var_opt with
    | Some (n,t) -> Some t
    | None ->
        if table.isRoot then
          if check_block_only then
            None
          else
            try_find_id (MethodInfo(table.Owner.Value.Value)) id check_block_only env
        else
          let parent = table.Parent.Value.Value
          try_find_id (BlockTable(parent)) id check_block_only env
  | MethodInfo(table) ->
      let method_param_opt = table.Params |> List.tryFind (fun (t,i) -> id = i)
      match method_param_opt with
      | Some (t,i) -> Some t
      | None -> try_find_id (ClassTable(table.OwnerClass.Value)) id check_block_only env
  | ClassTable(table) ->
      if table.Fields.ContainsKey(id) then
        Some (table.Fields.[id].Type)
      elif table.InheritedClass.IsSome then
        let inherited_class_table = env.SymbolTable.[table.InheritedClass.Value]
        try_find_id (ClassTable(inherited_class_table)) id check_block_only env
      else
        None

let try_find_class (table : Globals) (id : Id) : Option<SynType> =
  if table.ContainsKey(id) then
    let class_table = table.[id]
    Some (Identifier(class_table.Name))
  else
    None

let rec try_find_field (table : Globals) (class_id : Id) (var_id : Id) : Option<SynType> =
  if table.ContainsKey(class_id) then
    let class_table = table.[class_id]
    if class_table.Fields.ContainsKey(var_id) then
      let field_table = class_table.Fields.[var_id]
      Some (field_table.Type)
    elif class_table.InheritedClass.IsSome then
      try_find_field table class_table.InheritedClass.Value var_id
    else
      None
  else
    failwith "Bug in symbol table: expected class table cannot be found in try_find_field"

let rec try_find_method (table : Globals) (class_id : Id) (method_id : Id) : Option<MethodRetType * (Formals list)> =
  if table.ContainsKey(class_id) then
    let class_table = table.[class_id]
    if class_table.Methods.ContainsKey(method_id) then
      Some (class_table.Methods.[method_id].RetType,class_table.Methods.[method_id].Params)
    elif class_table.InheritedClass.IsSome then
      try_find_method table class_table.InheritedClass.Value method_id
    else
      None      
  else
    failwith "Bug in symbol table: expected class table cannot be found in try_find_method"

let rec is_super_class (table : Globals) (super : Id) (current_class : Id)  : bool =
  if table.ContainsKey(current_class) then
    let class_table = table.[current_class]
    match class_table.InheritedClass with
    | Some parent ->
        if parent = super then
          true
        else
          is_super_class table super parent
    | None -> false
  else
    failwith "Bug in symbol table: expected class table cannot be found in is_super_class"

let rec get_method_ret_type (block_table : BlockTable) : MethodRetType =
  if block_table.isRoot then
    match block_table.Owner with
    | Some owner ->
        owner.Value.RetType
    | None -> failwith "Bug in symbol table: expected an owner method for block in get_method_ret_type"
  else
    get_method_ret_type block_table.Parent.Value.Value

let get_method_info (table : Globals) (c_id : Id) (m_id : Id) : Ref<MethodInfo> =
  let current_class_table = table.[c_id]
  let method_info_opt = try_get (current_class_table.Methods) m_id
  match method_info_opt with
  | Some m_info -> ref m_info
  | None -> failwith "Bug in symbol table: could not find expected method in a class"

  
  
